10654. Unique color

 

A tree with n vertices numbered from 1 to n is given. The i-th edge of the tree connects vertices ai and bi​. Vertex i is colored with color ci (in this problem, colors are represented by integers).

A vertex x is considered good if the shortest path from vertex 1 to vertex x does not contain any other vertex with the same color as vertex x, except for x itself.

Find all good vertices.

 

Input. The first line contains an integer n (2 ≤ n ≤ 105) – the number of vertices.

The second line contains n integers – the colors of the vertices: c1, c2, ..., cn (1 ≤ ci ≤ 105).

Each of the following n – 1 lines contains two integers ai and bi (1 ≤ ai, bi ≤ n) – the edges of the tree.

 

Output. Print all good vertices, one per line, in ascending order of their indices.

 

Sample input 1

Sample output 1

6

2 7 1 8 2 8

1 2

3 6

3 2

4 3

2 5

1

2

3

4

6

 

 

Sample input 2

Sample output 2

10

3 1 4 1 5 9 2 6 5 3

1 2

2 3

3 4

4 5

5 6

6 7

7 8

8 9

9 10

1

2

3

5

6

7

8

 

 

SOLUTION

graphs – depth first search

 

Algorithm analysis

Start a depth-first search from vertex 1. When entering vertex v with color color[v], increment the value of used[color[v]] by 1. The value used[color[v]] represents how many times the color color[v] has appeared on the path from vertex 1 to vertex v (inclusive). If this color appears exactly once on the path, then vertex v is considered good, and we add it to the result set.

When exiting vertex v, decrement the value of used[color[v]] by 1.

 

Example

The graph from the first example looks as follows:

Vertex 5 is not good because on the path 1 – 2 – 5, vertices 5 and 1 have the same color.

Vertex 6 is good because on the path 1 – 2 – 3 – 6, the colors of all intermediate vertices differ from the color of vertex 6.

 

Algorithm implementation

Store the input graph in the adjacency list g. Declare the arrays.

 

vector<int> used, color;

vector<vector<int>> g;

set<int> st;

 

The dfs function implements depth-first search. The variable par is the parent of vertex v.

 

void dfs(int v, int par)

{

 

Vertex v has the color color[v]. Mark that a vertex with this color is visited on the path from vertex 1.

 

  used[color[v]]++;

 

The value used[color[v]] indicates how many times the color color[v] has appeared on the path from vertex 1 to vertex v (inclusive). If this color appears exactly once on the path, then vertex v is considered good, and we add it to the result set st.

 

  if (used[color[v]] == 1) st.insert(v);

 

Iterate over all edges (v, to) originating from vertex v. If vertex to is not the parent of v (to ≠ par), start a depth-first search from to. In this case, the parent of to becomes v.

 

  for (int to : g[v])

    if (to != par) dfs(to, v);

 

When exiting vertex v, decrement the value of used[color[v]] by 1.

 

  used[color[v]]--;

}

 

The main part of the program. Read the number of vertices n and the array of colors.

 

scanf("%d", &n);

color.resize(n + 1);

for (i = 1; i <= n; i++)

  scanf("%d", &color[i]);

 

Read the graph.

 

used.resize(100001);

g.resize(n + 1);

for (i = 1; i < n; i++)

{

  scanf("%d %d", &x, &y);

  g[x].push_back(y);

  g[y].push_back(x);

}

 

Start a depth-first search from vertex 1.

 

dfs(1, 1);

 

Print the good vertices. They are contained in the set st.

 

for (int val : st)

  printf("%d\n", val);